{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "multi_worker.ipynb",
      "version": "0.3.2",
      "provenance": [],
      "private_outputs": true,
      "collapsed_sections": [
        "Tce3stUlHN0L"
      ],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "cells": [
    {
      "metadata": {
        "colab_type": "text",
        "id": "Tce3stUlHN0L"
      },
      "cell_type": "markdown",
      "source": [
        "##### Copyright 2019 The TensorFlow Authors.\n",
        "\n"
      ]
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "tuOe1ymfHZPu",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "colab_type": "text",
        "id": "MfBg1C5NB3X0"
      },
      "cell_type": "markdown",
      "source": [
        "# Multi-worker Training in TensorFlow\n",
        "\n",
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/alpha/tutorials/distribute/multi_worker\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on TensorFlow.org</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/r2/tutorials/distribute/multi_worker.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/docs/blob/master/site/en/r2/tutorials/distribute/multi_worker.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "metadata": {
        "colab_type": "text",
        "id": "xHxb-dlhMIzW"
      },
      "cell_type": "markdown",
      "source": [
        "## Overview\n",
        "\n",
        "This tutorial demonstrates how `tf.distribute.Strategy` can be used for distributed multi-worker training with `tf.estimator`.  If you write your code using `tf.estimator`, and you're interested in scaling beyond a single machine with high performance, this tutorial is for you.\n",
        "\n",
        "Before getting started, please read the [`tf.distribute.Strategy` guide](../../guide/distribute_strategy.ipynb).  The [multi-GPU training tutorial](./keras.ipynb) is also relevant, because this tutorial uses the same model.\n",
        "\n"
      ]
    },
    {
      "metadata": {
        "colab_type": "text",
        "id": "MUXex9ctTuDB"
      },
      "cell_type": "markdown",
      "source": [
        "## Setup\n",
        "\n",
        "First, setup TensorFlow and the necessary imports."
      ]
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "IqR2PQG4ZaZ0",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "from __future__ import absolute_import, division, print_function"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "bnYxvfLD-LW-",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "!pip install tensorflow-gpu==2.0.0-alpha0\n",
        "import tensorflow_datasets as tfds\n",
        "import tensorflow as tf\n",
        "\n",
        "import os, json"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "hPBuZUNSZmrQ",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## Input Function\n",
        "\n",
        "This tutorial uses the MNIST dataset from [TensorFlow Datasets](https://www.tensorflow.org/datasets).  The code here is similar to the [multi-GPU training tutorial](./keras.ipynb) with one key difference: when using Estimator for multi-worker training, it is necessary to shard the dataset by the number of workers to ensure model convergence.  The input data is sharded by worker index, so that each worker processes 1/`num_workers` distinct portions of the dataset."
      ]
    },
    {
      "metadata": {
        "id": "dma_wUAxZqo2",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "BUFFER_SIZE = 10000\n",
        "BATCH_SIZE = 64\n",
        "\n",
        "def input_fn(mode, input_context=None):\n",
        "  datasets, ds_info = tfds.load(name='mnist',\n",
        "                                with_info=True,\n",
        "                                as_supervised=True)\n",
        "  mnist_dataset = (datasets['train'] if mode == tf.estimator.ModeKeys.TRAIN else\n",
        "                   datasets['test'])\n",
        "\n",
        "  def scale(image, label):\n",
        "    image = tf.cast(image, tf.float32)\n",
        "    image /= 255\n",
        "    return image, label\n",
        "\n",
        "  if input_context:\n",
        "    mnist_dataset = mnist_dataset.apply(tf.data.experimental.filter_for_shard(\n",
        "        input_context.num_input_pipelines, input_context.input_pipeline_id))\n",
        "  return mnist_dataset.map(scale).shuffle(BUFFER_SIZE).batch(BATCH_SIZE)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "4BlcVXMhB59T",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "Another reasonable approach to achieve convergence would be to shuffle the dataset with distinct seeds at each worker."
      ]
    },
    {
      "metadata": {
        "id": "8YFpxrcsZ2xG",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## Multi-worker Configuration\n",
        "\n",
        "One of the key differences in this tutorial, compared to multi-GPU training, is the multi-worker setup.  The `TF_CONFIG` environment variable is the standard way to specify the cluster configuration to each worker that is part of the cluster.\n",
        "\n",
        "There are two components of `TF_CONFIG`: `cluster` and `task`.  `cluster` provides information about the entire cluster, namely the workers and parameter servers in the cluster.  `task` provides information of the current task.  In this example, the task `type` is `worker` and the task `index` is `0`. \n",
        "\n",
        "For illustration purposes, this tutorial shows how one may set a `TF_CONFIG` with a single worker on `localhost`.  In practice, users would create multiple workers on an external IP address and port, and set `TF_CONFIG` on each worker appropriately, i.e. modify the task `index`.\n",
        "\n",
        "Warning: Do not execute the following code in Colab.  TensorFlow's runtime will attempt to create a gRPC server at the specified IP address and port, which will likely fail.\n",
        "\n",
        "```\n",
        "NUM_WORKERS = 1\n",
        "IP_ADDRS = ['localhost']\n",
        "PORTS = [12345]\n",
        "\n",
        "os.environ['TF_CONFIG'] = json.dumps({\n",
        "    'cluster': {\n",
        "        'worker': ['%s:%d' % (IP_ADDRS[w], PORTS[w]) for w in range(NUM_WORKERS)]\n",
        "    },\n",
        "    'task': {'type': 'worker', 'index': 0}\n",
        "})\n",
        "```\n",
        "\n"
      ]
    },
    {
      "metadata": {
        "id": "qDreJzTffAP5",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## Define the model\n",
        "\n",
        "Write the layers, the optimizer, and the loss function for training.  This tutorial defines the model with Keras layers, similar to the [multi-GPU training tutorial](./keras.ipynb)."
      ]
    },
    {
      "metadata": {
        "id": "WNvOn_OeiUYC",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "LEARNING_RATE = 1e-4\n",
        "def model_fn(features, labels, mode):\n",
        "  model = tf.keras.Sequential([\n",
        "      tf.keras.layers.Conv2D(32, 3, activation='relu', input_shape=(28, 28, 1)),\n",
        "      tf.keras.layers.MaxPooling2D(),\n",
        "      tf.keras.layers.Flatten(),\n",
        "      tf.keras.layers.Dense(64, activation='relu'),\n",
        "      tf.keras.layers.Dense(10, activation='softmax')\n",
        "  ])\n",
        "  logits = model(features, training=False)\n",
        "  \n",
        "  if mode == tf.estimator.ModeKeys.PREDICT:\n",
        "    predictions = {'logits': logits}\n",
        "    return tf.estimator.EstimatorSpec(labels=labels, predictions=predictions)\n",
        "  \n",
        "  optimizer = tf.compat.v1.train.GradientDescentOptimizer(\n",
        "      learning_rate=LEARNING_RATE)\n",
        "  loss = tf.keras.losses.SparseCategoricalCrossentropy(\n",
        "      from_logits=True)(labels, logits)\n",
        "  if mode == tf.estimator.ModeKeys.EVAL:\n",
        "    return tf.estimator.EstimatorSpec(mode, loss=loss)\n",
        "  \n",
        "  return tf.estimator.EstimatorSpec(\n",
        "      mode=mode,\n",
        "      loss=loss,\n",
        "      train_op=optimizer.minimize(\n",
        "          loss, tf.compat.v1.train.get_or_create_global_step()))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "P94PrIW_kSCE",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "Note that while the learning rate is fixed in this example, in general it may be necessary to adjust the learning rate based on the global batch size."
      ]
    },
    {
      "metadata": {
        "colab_type": "text",
        "id": "UhNtHfuxCGVy"
      },
      "cell_type": "markdown",
      "source": [
        "## MultiWorkerMirroredStrategy\n",
        "\n",
        "To train the model, use an instance of `tf.distribute.experimental.MultiWorkerMirroredStrategy`.  `MultiWorkerMirroredStrategy` creates copies of all variables in the model's layers on each device across all workers.  It uses `CollectiveOps`, a TensorFlow op for collective communication, to aggregate gradients and keep the variables in sync.  The [`tf.distribute.Strategy` guide](../guide/distribute_strategy.ipynb) has more details about this strategy."
      ]
    },
    {
      "metadata": {
        "id": "1uFSHCJXMrQ-",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "FMy2VM4Akzpr",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "`MultiWorkerMirroredStrategy` provides multiple implementations via the [`CollectiveCommunication`](https://github.com/tensorflow/tensorflow/blob/a385a286a930601211d78530734368ccb415bee4/tensorflow/python/distribute/cross_device_ops.py#L928) parameter.  `RING` implements ring-based collectives using gRPC as the cross-host communication layer.  `NCCL` uses [Nvidia's NCCL](https://developer.nvidia.com/nccl) to implement collectives.  `AUTO` defers the choice to the runtime.  The best choice of collective implementation depends upon the number and kind of GPUs, and the network interconnect in the cluster."
      ]
    },
    {
      "metadata": {
        "id": "H47DDcOgfzm7",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## Train and evaluate the model\n",
        "\n",
        "Next, specify the distribution strategy in the `RunConfig` for the estimator, and train and evaluate by invoking `tf.estimator.train_and_evaluate`.  This tutorial distributes only the training by specifying the strategy via `train_distribute`.  It is also possible to distribute the evaluation via `eval_distribute`."
      ]
    },
    {
      "metadata": {
        "id": "BcsuBYrpgnlS",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "config = tf.estimator.RunConfig(train_distribute=strategy)\n",
        "\n",
        "classifier = tf.estimator.Estimator(\n",
        "    model_fn=model_fn, model_dir='/tmp/multiworker', config=config)\n",
        "tf.estimator.train_and_evaluate(\n",
        "    classifier,\n",
        "    train_spec=tf.estimator.TrainSpec(input_fn=input_fn),\n",
        "    eval_spec=tf.estimator.EvalSpec(input_fn=input_fn)\n",
        ")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "XVk4ftYx6JAO",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## Performance\n",
        "\n",
        "You now have a model and a multi-worker capable Estimator powered by `tf.distribute.Strategy`.  You can try the following techniques to tweak performance of multi-worker training.\n",
        "\n",
        "*   `MultiWorkerMirroredStrategy` provides multiple [collective communication implementations](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/distribute/cross_device_ops.py).  `RING` implements ring-based collectives using gRPC as the cross-host communication layer.  `NCCL` uses [Nvidia's NCCL](https://developer.nvidia.com/nccl) to implement collectives.  `AUTO` defers the choice to the runtime.  The best choice of collective implementation depends upon the number and kind of GPUs, and the network interconnect in the cluster.  To override the automatic choice, specify a valid value to the `communication` parameter of `MultiWorkerMirroredStrategy`'s constructor, e.g. `communication=tf.distribute.experimental.CollectiveCommunication.NCCL`.\n",
        "*   The batch size specified here is per-GPU.  In general, the largest batch size that fits the GPU memory is advisable.\n",
        "*    Cast the variables to `tf.float` if possible.  The official ResNet model includes [an example](https://github.com/tensorflow/models/blob/8367cf6dabe11adf7628541706b660821f397dce/official/resnet/resnet_model.py#L466) of how this can be done.\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "metadata": {
        "id": "AW0Hb2xM6EGX",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## Other Code Examples\n",
        "\n",
        "1.   [End to end example](https://github.com/tensorflow/ecosystem/tree/master/distribution_strategy) for multi worker training in tensorflow/ecosystem using Kubernetes templates. This example starts with a Keras model and converts it to an Estimator using the `tf.keras.estimator.model_to_estimator` API. \n",
        "2.   Official [ResNet50](https://github.com/tensorflow/models/blob/master/official/resnet/imagenet_main.py) model, which can be trained using either `MirroredStrategy` or `MultiWorkerMirroredStrategy`.\n",
        "\n"
      ]
    }
  ]
}
